import modal

app = modal.App("lightfm-openmp-benchmark")

# Create image with LightFM compiled with OpenMP
image = (
    modal.Image.debian_slim(python_version="3.10")
    .apt_install("gcc", "g++", "libgomp1", "git")
    ## begin[broken_parallelism]
    # .pip_install("numpy", "scipy", "requests", "scikit-learn", "cython")
    # .run_commands("git clone https://github.com/daviddavo/lightfm.git /tmp/lightfm")
    #  .run_commands("cd /tmp/lightfm && python cythonize setup.py && pip install .")
    ## end[broken_parallelism
    .pip_install("numpy", "scipy", "requests", "scikit-learn", "cython", "lightfm")
)


@app.function(image=image, cpu=4)
def benchmark():
    import numpy as np
    from scipy.sparse import coo_matrix
    from lightfm import LightFM
    import time
    import os
    import sys

    print(f"Platform: {sys.platform}")
    print(f"CPU count: {os.cpu_count()}")

    # Check which module is loaded
    try:
        import lightfm._lightfm_fast_openmp

        print("✓ Using OpenMP version")
    except ImportError as e:
        print(f"✗ OpenMP version not available: {e}")

    try:
        import lightfm._lightfm_fast_no_openmp

        print("✓ Using no-OpenMP version")
    except ImportError as e:
        print(f"✗ No-OpenMP version not available: {e}")

    # Create synthetic dataset
    n_users = 30_000
    n_items = 30_000
    n_interactions = 300_000

    users = np.random.randint(0, n_users, n_interactions)
    items = np.random.randint(0, n_items, n_interactions)
    data = np.ones(n_interactions)

    interactions = coo_matrix((data, (users, items)), shape=(n_users, n_items))

    # Benchmark
    print("\nBenchmark results:")
    baseline_time = None
    for threads in [1, 2, 4]:
        model = LightFM(no_components=128, loss="warp")
        start = time.time()
        model.fit(interactions, epochs=10, num_threads=threads)
        elapsed = time.time() - start

        if threads == 1:
            baseline_time = elapsed

        speedup = baseline_time / elapsed if baseline_time else 1.0
        print(f"{threads} threads: {elapsed:.2f}s ({speedup:.2f}x speedup)")

    return "Benchmark complete"


@app.local_entrypoint()
def main():
    result = benchmark.remote()
    print(result)